// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

#pragma once

#include <cstdlib>
#include <deque>
#include <functional>
#include <iostream>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <vector>

#include "coin2/base/config/app_config.h"
#include "coin2/base/log.h"
#include "coin2/base/time.h"
#include "coin2/exchange/base/api_base/raw_subscriber.h"
#include "coin2/exchange/base/executor/live_executor.h"
#include "coin2/exchange/base/feed/latency_manager.h"
#include "coin2/exchange/base/feed/ws/feed_subscription_spec.h"
#include "coin2/exchange/base/feed/ws/request_generator.h"
#include "coin2/exchange/base/feed/ws/ws_connection.h"
#include "coin2/exchange/base/market/market_exchange_api.h"
#include "coin2/exchange/factory/feed/request_generator_factory.h"
#include "coin2/exchange/feed.pb.h"
#include "coin2/exchange/util/rate_limit.h"
#include "coin2/flow/symbol_group.h"
#include "coin2/flow/topic_string.h"

namespace coin2::exchange::base::feed {

using coin2::exchange::feed::WebsocketConfig;
using coin2::exchange::base::symbology::IProduct;
using coin2::exchange::base::api_base::HealthReport;

inline void ValidateWebsocketConfig(
    const WebsocketConfig& config,
    const FeedSubscriptionSpec& spec) {
  CHECK_THROW(config.has_num_workers()) << "[Config] live.num_workers missing";
  if (config.has_worker_id()) {
    CHECK_EQ_THROW(config.num_workers(), 1) << "[Config] live.num_workers must be 1, "
                                               "if worker_id specified.";
    CHECK_THROW((config.worker_id() == 1) || (config.worker_id() == 2))
        << "[Config] live.worker_id must be 1 or 2.";
  }

  // do channels really exist?
  for (const auto& ch : config.channels()) {
    CHECK_THROW(spec.channels.find(ch) != spec.channels.end())
        << "[Config] channel " << ch << " does not exist in " << spec.mea.String();
  }
}

class WsSubscriber : public std::enable_shared_from_this<WsSubscriber>,
                     public api_base::BaseRawAsyncSubscriber {
 public:
  using ConnectionKey = std::pair<int, int>;

  WsSubscriber(
      const FeedSubscriptionSpec& spec,
      const WebsocketConfig& config,
      const std::vector<WsConnectionConfig>& connection_configs,
      api_base::LiveAsyncExecutor* flow_sub,
      bool enable_telemetry)
      : BaseRawAsyncSubscriber(flow_sub),
        spec_(spec),
        config_(config),
        connection_configs_(connection_configs),
        strand_(*ioc()),
        subscribe_timer_(*ioc()),
        request_generator_(factory::feed::NewRequestGenerator(spec.mea, &config)),
        scheduler_(new Scheduler(&strand_)),
        latency_manager_(spec.mea, config, ioc()),
        enable_telemetry_(enable_telemetry) {
    ValidateWebsocketConfig(config, spec);

    if (spec_.unauthorized_limit.connection) {
      unauthorized_connection_limit_ = std::make_unique<WindowLimiter>(
          spec_.unauthorized_limit.connection->first,
          spec_.unauthorized_limit.connection->second);
    }

    if (spec_.reconnect_delay_ns > 0 ||
        spec_.periodic_force_reconnect_minute > 0 ||
        config_.reconnect_slowest() ||
        config_.monitor_connection()) {
      scheduler_->Start();
    }
  }

  void ForceReconnectWorker(const WsConnectionConfig* config) {
    AddForSymbolGroupAndWorker(config);

    if (config->mea_config->unauthorized_limit.max_subscription) {
      CHECK_LE_THROW(config->symbols.size(),
          config->mea_config->unauthorized_limit.max_subscription.value()) <<
          "[WsSubscriber] exceeds unauthorized max_subscription for " <<
          config->mea_config->mea.String();
    }
  }

  void RegisterPubTopic() override { flow_sub_->distributor()->Publish(spec_.mea.SnakeString()); }

  int GetWorkerId(int worker_index) {
    if (config_.has_worker_id()) {
      return config_.worker_id() + worker_index - 1;
    } else {
      return worker_index;
    }
  }

  void AddForSymbolGroupAndWorker(const WsConnectionConfig* config) {
    std::scoped_lock _(connections_mutex_);
    ConnectionKey key(config->connection_id, config->worker_id);
    if (connections_.count(key) > 0) {
      auto conn = connections_.at(key);
      conn->PostCleanUp();
    }
    connections_.erase(key);
    if (std::find(pending_connections_.begin(),
                  pending_connections_.end(),
                  config)
        == pending_connections_.end()) {
      pending_connections_.push_back(config);
    }
  }

  void UnsubscribeForSymbolGroupAndWorker(
      const WsConnectionConfig* config,
      const std::vector<std::string>& removed_symbols) {
    std::scoped_lock _(connections_mutex_);
    ConnectionKey key(config->connection_id, config->worker_id);
    if (connections_.count(key) > 0) {
      auto conn = connections_.at(key);
      conn->PostUnsubcribe(removed_symbols);
    }
  }

  void RemoveForSymbolGroupAndWorker(const WsConnectionConfig* config) {
    std::scoped_lock _(connections_mutex_);
    ConnectionKey key(config->connection_id, config->worker_id);
    if (connections_.count(key) > 0) {
      auto conn = connections_.at(key);
      conn->PostCleanUp();
    }
    connections_.erase(key);
    auto iter = std::find(pending_connections_.begin(),
                            pending_connections_.end(),
                            config);
    if (iter != pending_connections_.end()) {
      pending_connections_.erase(iter);
    }
  }

  bool ExtractForSymbolGroupAndWorker(const WsConnectionConfig** config) {
    std::scoped_lock _(connections_mutex_);
    if (!pending_connections_.empty()) {
      bool add_now = true;
      if (unauthorized_connection_limit_) {
        add_now = unauthorized_connection_limit_->CheckAndSetWithTimestamp(GetCurrentTimestamp());
      }

      if (add_now) {
        *config = pending_connections_.front();
        pending_connections_.pop_front();
        return true;
      }
    }

    return false;
  }

  // needs to be thread-safe!
  // network disconnects
  // => request for connect a1, a2, b1, b2, ..
  // => multiple ConnectForSymbolGroupAndWorker(..) at the same time
  virtual void ConnectForSymbolGroupAndWorker(const WsConnectionConfig* connection_config) {
    std::weak_ptr<WsSubscriber> self(shared_from_this());
    std::shared_ptr<LatencyExtractor> extractor(nullptr);
    if (config_.reconnect_slowest() || config_.monitor_connection()) {
      extractor = latency_manager_.AddLatencyExtractor(connection_config);
    }
    std::shared_ptr<WsConnection> conn = std::make_shared<WsConnection>(
        *connection_config,
        [self, connection_config] {
          if (auto s = self.lock()) {
            boost::asio::post(s->strand_, [self, connection_config] {
              if (auto s = self.lock()) {
                s->latency_manager_.RemoveLatencyExtractor(connection_config);
                s->AddForSymbolGroupAndWorker(connection_config);
              } else {
                LOG(ERROR) << "boost::asio handler object (WsSubscriber) does not exist";
              }
            });
          } else {
            LOG(ERROR) << "boost::asio handler object (WsSubscriber) does not exist";
          }
        },
        flow_sub_,
        request_generator_.get(),
        extractor,
        enable_telemetry_);
    conn->RunAsync();
    {
      std::scoped_lock _(connections_mutex_);
      ConnectionKey k(connection_config->connection_id, connection_config->worker_id);
      connections_.insert({k, conn});
    }
  }

  void SubscribeAsync() override;

  void Subscribe(boost::asio::yield_context yield) {
    while (sub_running_) {
      const WsConnectionConfig* connection_config;

      if (ExtractForSymbolGroupAndWorker(&connection_config)) {
        if (spec_.reconnect_delay_ns > 0) {
          std::string task_name = std::to_string(connection_config->connection_id) +
                                  std::to_string(connection_config->worker_id);
          std::weak_ptr<WsSubscriber> self(shared_from_this());
          scheduler_->CallLater(
              [self, connection_config] {
                if (auto s = self.lock()) {
                  s->ConnectForSymbolGroupAndWorker(connection_config);
                } else {
                  LOG(ERROR) << "boost::asio handler object (WsSubscriber) does not exist";
                }
              },
              spec_.reconnect_delay_ns,
              task_name);
        } else {
          ConnectForSymbolGroupAndWorker(connection_config);
        }
      }

      subscribe_timer_.expires_from_now(std::chrono::milliseconds(100));
      subscribe_timer_.async_wait(yield);
    }
  }

  void ForceReconnect() {
    pending_connections_.clear();
    for (auto& it : connection_configs_) {
      AddForSymbolGroupAndWorker(&it);

      if (it.mea_config->unauthorized_limit.max_subscription) {
        CHECK_LE_THROW(it.symbols.size(),
            it.mea_config->unauthorized_limit.max_subscription.value()) <<
            "[WsSubscriber] exceeds unauthorized max_subscription for " <<
            it.mea_config->mea.String();
      }
    }
  }

  void GetHealthReport(int indent, HealthReport* health) const override {
    std::map<ConnectionKey, std::string> infos_status;
    std::map<ConnectionKey, std::string> infos_last;
    if (latency_manager_.Enabled()) {
      auto packets = latency_manager_.CollectConnectionInfo();
      for (auto const& [k, packet] : packets) {
        std::stringstream ss;
        if (packet) {
          ss << "last_update=" 
            << DebugStringFromDuration(GetCurrentTimestamp() - packet->timestamp);
          infos_last[k] = ss.str();
        } else {
          // conection start but no data received
          infos_last[k] = "last_update= ";
        }
      }
    }
    {
      std::scoped_lock _(connections_mutex_);
      health->connection_num += connections_.size();
      for (auto const& [k, _] : connections_) {
        infos_status[k] = "STARTING";
      }

      for (const auto& config : pending_connections_) {
        ConnectionKey key(config->connection_id, config->worker_id);
        infos_status[key] = "PENDING";
      }
    }

    for (const auto& it : connection_configs_) {
      for (int ind = 0; ind < indent * 4; ind++) {
        health->report += ' ';
      }

      ConnectionKey key(it.connection_id, it.worker_id);

      std::string name = it.topic + "_" + it.url.name + "_" + std::to_string(it.connection_id);
      std::string health_status = "CLOSED";
      auto iter_status = infos_status.find(key);
      if (iter_status != infos_status.end()) {
        health_status = iter_status->second;
      }
      std::string last;
      auto iter_last = infos_last.find(key);
      if (iter_last != infos_last.end()) {
        last = iter_last->second;
      }

      health->report += name;
      health->report += " (";
      health->report += health_status;
      health->symbol_num += it.symbols.size();
      health->report += ", symbols=" + std::to_string(it.symbols.size());
      if (!last.empty()) {
        health->report += ", " + last;
      }
      health->report += ")\n";
    }
  }

  void CleanUp() override {
    scheduler_->Stop();
  }

  bool BanProduct(const std::vector<const IProduct*>& products) override;

 protected:
  void FindAndRemoveSymbolFromConfig(
      const std::vector<const IProduct*> products,
      std::map<WsConnectionConfig*, std::vector<std::string>>* connection_configs);

  const FeedSubscriptionSpec& spec_;
  const WebsocketConfig& config_;
  // connection per symbol group config
  std::vector<WsConnectionConfig> connection_configs_;
  std::deque<const WsConnectionConfig*> pending_connections_;
  std::map<ConnectionKey, std::shared_ptr<WsConnection>> connections_;
  mutable std::mutex connections_mutex_;
  std::unique_ptr<WindowLimiter> unauthorized_connection_limit_;

 public:
  boost::asio::io_context::strand strand_;

 protected:
  boost::asio::steady_timer subscribe_timer_;
  std::unique_ptr<RequestGenerator> request_generator_;
  std::shared_ptr<Scheduler> scheduler_;

 public:
  LatencyManager latency_manager_;

 protected:
  bool enable_telemetry_;
};

}  // namespace coin2::exchange::base::feed
